package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;


import nlp.langmodel.LanguageModel;
import nlp.util.Counter;

public class InterpolatedTrigramModel implements LanguageModel {
	private Collection<List<String>> trainingCollection;
	private Collection<List<String>> validationCollection;
	private UnigramModel unigramModel;
	private BigramModel bigramModel;
	private TrigramModel trigramModel;
	private Trigram<Double, Double, Double> lambdas = null;
	private static final String START = "<S>";
	private static final String STOP = "</S>";
	
	public InterpolatedTrigramModel(Collection<List<String>> trainingCollection, Collection<List<String>> validationCollection) {
		this.trainingCollection = trainingCollection;
		this.validationCollection = validationCollection;
		unigramModel = new UnigramModel(this.trainingCollection, this.validationCollection);
		bigramModel = new BigramModel(this.trainingCollection, this.validationCollection);
		trigramModel = new TrigramModel(this.trainingCollection, this.validationCollection);
		lambdas = resolveLambda();
//		lambdas = new Trigram<Double, Double, Double>(.3, .55, .15); // (0.30000000000000004, 0.55, 0.1499999999999999)
		
		System.out.println("lambdaPair is " + lambdas);
		System.out.println("LanguageModelTester.calculatePerplexity() is " + LanguageModelTester.calculatePerplexity(this, validationCollection));
		
	}

	
	Trigram<Double, Double, Double> resolveLambda() {
		double gap = .01;
		int times = (int)(1.0 / gap);
		
		double lambda1 = 0.0;
		double lambda2 = 0.0;
		double lambda3 = 1.0 - lambda1 - lambda2;
		double perplexity = Double.POSITIVE_INFINITY;
		int loop = 0;
		for (int i = 0; i < times; i++) {
			for (int j = 0; j < times - i; j++) {
				double temp = calculatePerplexity(validationCollection, i * gap, j * gap, 1.0 - (i + j) * gap);
				if (temp < perplexity) {
					perplexity = temp;
					lambda1 = gap * i;
					lambda2 = gap * j;
					lambda3 = 1.0 - lambda1 - lambda2;
				}
				System.out.println("temp perplexity of round " + (loop++) + " is " + temp);
			}
		}
		System.out.println("\nlowest perplexity " + perplexity);
		return new Trigram<Double, Double, Double>(lambda1, lambda2, lambda3);
	}

	double getInterpolatedWordLogProbability(List<String> sentence, int index, double lambda1, double lambda2, double lambda3) {
		if (index <= 1) {
			return 0.0;
		}
		
		double pUnigram = unigramModel.getWordLogProbability(sentence, index);
		double pBigram = bigramModel.getWordLogProbability(sentence, index);
		double pTrigram = trigramModel.getWordLogProbability(sentence, index);
//		System.out.println("pUnigram is " + pUnigram + ", and pBigram is " + pBigram);
		double temp = Math.pow(2.0, pUnigram) * lambda1 + Math.pow(2.0, pBigram) * lambda2 + Math.pow(2.0, pTrigram) * lambda3;
		return Math.log(temp) / Math.log(2.0);
	}
	
	double getInterplatedSentenceLogProbability(List<String> sentence, double lambda1, double lambda2, double lambda3) {
		double logProbability = 0.0;
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		for (int i = 0; i < stoppedSentence.size(); i++) {
			logProbability += getInterpolatedWordLogProbability(stoppedSentence, i, lambda1, lambda2, lambda3);
		}
//		System.out.println("sentence p is " + logProbability);
		return logProbability;
	}
	
	private double calculatePerplexity(Collection<List<String>> sentenceCollection, double lambda1, double lambda2, double lambda3) {
	    double logProbability = 0.0;
	    double numSymbols = 0.0;
	    for (List<String> sentence : sentenceCollection) {
	      logProbability += getInterplatedSentenceLogProbability(sentence, lambda1, lambda2, lambda3);
	      numSymbols += sentence.size();
	    }
	    double avgLogProbability = logProbability / numSymbols;
	    double perplexity = Math.pow(.5, avgLogProbability);
	    return perplexity;
	  }
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		double probability = 0.0;
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(0, START);
		stoppedSentence.add(0, START);
		stoppedSentence.add(STOP);
		
		double l1 = lambdas.first;
		double l2 = lambdas.second;
		double l3 = lambdas.third;
		for (int index = 0; index < stoppedSentence.size(); index++) {
			probability += getInterpolatedWordLogProbability(stoppedSentence, index, l1, l2, l3);
		}
//		System.out.println("sentence probability is " + Math.pow(2.0, probability));
		return Math.pow(2.0, probability);
	}

	String generateWord(Bigram<String, String> history) {
		Counter<String> counter = unigramModel.getUnigramCounter();
		double sample = Math.random();
		double sum = 0.0;
		String second = history.second;
		
		for (String word : counter.keySet()) {
			double pUni = Math.pow(2.0, unigramModel.getPWord(word));
			double pBi = Math.pow(2.0, bigramModel.getPWord(second, word));
			double pTri = Math.pow(2.0, trigramModel.getPWord(history, word));
			sum +=  lambdas.getFirst() * pUni + lambdas.getSecond() * pBi + lambdas.getThird() * pTri;
//			System.out.println("sum is " + sum + ", and sample is " + sample);
			if (sum > sample) {
				return word;
			}
		}
//		System.out.println("sum is " + sum);
//		return "*UNKNOWN*";
		return STOP;
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		String hist = START;
		Bigram<String, String> bigram = new Bigram<String, String>(hist, hist);
		String word = generateWord(bigram);
		while (!word.equals(STOP)) {
			sentence.add(word);
			bigram = new Bigram<String, String>(hist, word);
			hist = word;
			word = generateWord(bigram);
		}
		return sentence;
	}
}
